#!/bin/bash

cd /opt/pwn.college

if [ -n "$1" ]
then
    ACTION="$1"
    shift
else
    ACTION="help"
fi

DOCKER_ARGS=${DOCKER_ARGS:--i}
[ -t 0 ] && DOCKER_ARGS="-t $DOCKER_ARGS"

DOJO_DIR=/opt/pwn.college
if [ -f $DOJO_DIR/data/config.env ]; then
    . $DOJO_DIR/data/config.env
fi

case "$ACTION" in
    # HELP: update: update dojo files (warning: does `git pull`), rebuild containers, and restart any changed services
    "update")
        git pull
        dojo sync
        dojo compose up -d --build
        ;;

    # HELP: sync: sync changed dojo files
    "sync")
        cp -rv --preserve=timestamps ctfd/* /opt/CTFd/
        cp -rv --preserve=timestamps dojo_plugin /opt/CTFd/CTFd/plugins/
        cp -rv --preserve=timestamps dojo_theme /opt/CTFd/CTFd/themes/
        ;;

    # HELP: enter [ -s ] USER_ID: enter a user's running container. -s for entering as root
    "enter")
        USER_SWITCH=""
        if [ "$1" == "-s" ]
        then
            USER_SWITCH="--user=root"
            shift
        fi
        DOJO_UID="$1"
        [ -n "${DOJO_UID//[0-9]}" ] && DOJO_UID=$(
            echo "select id from users where name='$DOJO_UID'" |
            $0 db -s
        )
        CONTAINER="user_$DOJO_UID"
        shift
        docker exec $DOCKER_ARGS $USER_SWITCH "$CONTAINER" bash
        ;;

    # HELP: compose ARGS: run a docker compose command with the config.env file loaded
    "compose")
        docker compose --env-file=/opt/pwn.college/data/config.env "$@"
        ;;

    # HELP: flask: drop into a flask shell in the ctfd container
    "flask")
        docker exec $DOCKER_ARGS ctfd flask shell "$@"
        ;;

    # HELP: db: launch a mysql client session, connected to the ctfd db
    "db")
        docker exec $DOCKER_ARGS db mysql -pctfd -Dctfd ctfd "$@"
        ;;

    # HELP: backup: does a dojo db backup into the `data/backups` directory.
    "backup")
        mkdir -p data/backups
        docker exec db mysqldump -pctfd --single-transaction --routines --triggers ctfd | gzip > "data/backups/db-$(date -Iseconds).sql.gz"
        ;;

    # HELP: restore PATH: restores a dojo db backup. Path arg is relative to the `data/backups` directory
    "restore")
        BACKUP_PATH="data/backups/$1"
        if [ -f "$BACKUP_PATH" ]; then
            gunzip < "$BACKUP_PATH" | docker exec -i db mysql -pctfd -Dctfd
        else
            echo "Error: missing file to restore from" >&2
        fi
        ;;

    # HELP: logs: display dojo logs
    "logs")
        journalctl -u pwn.college -f
        ;;

    # HELP: wait: wait for the dojo to finish starting
    "wait")
        dojo logs | while IFS= read -r line; do
            echo "$line"
            case "$line" in
                *"Finished pwn.college docker compose service"*)
                    exit 0
                    ;;
                *"Failed to start pwn.college docker compose service."*)
                    exit 1
                    ;;
            esac
        done
        ;;

    # HELP: start: start the dojo. Main entry point for the docker image.
    "start")
        dojo sync
        echo "[+] Initializing dojo."
        dojo-init
        echo "[+] Starting systemd. If this is the first start, the challenge container"
        echo "    will be built, which will take a very long time."
        exec /usr/bin/systemd
        ;;

    # HELP: help: displays this help
    "help")
        echo "Usage: $0 COMMAND [ ARG ]"
	echo
	echo "COMMANDS:"
	echo
	cat "$0" | grep "[H]ELP" | sed -e "s/.*H[E]LP: /\t/"
        ;;

    *)
        echo "Unknown command."
	echo ""
	"$0" help
        exit 1
        ;;
esac
